import os
from ultralytics import YOLO
import torch.nn as nn
import torch
import numpy as np
import cv2
import time

model = YOLO('./model/SpikeYOLO/weight/yolov8l.pt')

# 文件夹路径
folder_path = '../../data/Hyper'
folder_save = '../../data/Hyper_output'

# 记录已处理的文件
processed_files = set()

def to_rgb_image(hsi_img):
    """
    Convert a 30-band hyperspectral image to a false-color RGB image.

    hsi_img: ndarray, shape (H, W, 30)
    """
    # 通道顺序 [R, G, B]，对应 index [10, 5, 3]
    img_rgb = hsi_img[:, :, [23, 5, 3]]
    # 归一化到 0-255
    img_rgb = (img_rgb - img_rgb.min()) / (img_rgb.max() - img_rgb.min() + 1e-6)
    img_rgb = (img_rgb * 255).astype(np.uint8)
    return img_rgb

while True:
    # 获取文件夹中的所有文件
    files = os.listdir(folder_path)
    
    # 遍历文件夹中的文件
    for file in files:
        file_path = os.path.join(folder_path, file)
        
        # 检查是否是新文件
        if file not in processed_files and os.path.isfile(file_path):
            print(f"Processing new file: {file}")
            
            # 加载图像
            data = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
            if data is None:
                print(f"Failed to load image: {file}")
                continue
            
            # 创建 cube 并进行处理
            cube = np.ones((256, 512, 16))
            for j in range(16):
                rowStart = int(j / 4)
                colStart = int(j % 4)
                cube[:, :, j] = data[rowStart:1024:4, colStart:2048:4]
            
            # 将文件标记为已处理
            processed_files.add(file)

            # 将 cube 抽取三个波段转换成 tensor
            
            rgb_img = to_rgb_image(cube)
            img_tensor = torch.from_numpy(rgb_img).permute(2, 0, 1).float()  # HWC → CHW
            img_tensor /= 255.0  # 归一化到 [0, 1]
            # 添加 batch 维度
            img_tensor = img_tensor.unsqueeze(0)

            # 通过file_path生成保存路径
            save_path = os.path.join(folder_save, f"{os.path.splitext(file)[0]}_o.jpg")

            # 进行推理
            results = model(rgb_img, device=0)
            results[0].show()         # 显示检测结果（默认用 OpenCV 窗口）
            results[0].save(filename='result.jpg')  # 保存结果图像
    # 等待一段时间后再次检查文件夹
    time.sleep(1)
